#!/usr/bin/python/ 

""" 
This script takes output from
<timestamp>/pickled_output/ and bunches it into easy-to-handle arrays
so that we can make histograms, correlation plots, etc. 
""" 

from sys import argv
import os, re
import cPickle as pickle
from utils import is_timestamped
import numpy as np

SOMETIMES_KEYS = ['h_offset', 'h_peak', 'stdev_mul', 'sigma']

#Convenience function stolen from stackoverflow.com/questions/1548704
def purge(dire, pattern):
    for f in os.listdir(dire):
        if re.search(pattern, f):
            os.remove(os.path.join(dire, f))

def bunch_output(timestamp):
    keys = ['s', 's_filtered', 'overlap_p', 'overlap_m']

    durr = timestamp + '/pickled_output/'
    prfx = 'trajectory_data'    
    sffx = '.pkl'

    output_dict = {key:[] for key in keys}
    #Python ranges from 0 to n-1, 
    # bsub ranges from 1 to n, auto-adjust loop variable
    flnm_lst = os.listdir(durr)    
       
    #n_files = len(os.listdir(timestamp+'/pickled_output/'))
    #print n_files
    try:
        for flnm in flnm_lst:
            dict_idx={}
            with open(durr + flnm, 'r') as phil:
                dict_idx = pickle.load(phil)
            #print dict_idx['s']
            for key in keys:
                temp_val = dict_idx[key]
                output_dict[key].append(temp_val)
    except ValueError:
        pass #dirt
    
    with open(timestamp + '/output_arrays/output.pkl','w') as phil:
        pickle.dump(output_dict, phil)

    err = 0    
    return err

def rm_pickles(timestamp):
    output_contents = os.listdir(timestamp+'/output_arrays')   
    if 'output.pkl' not in output_contents:
        print "This output has not been made redundant, don't remove it."
        return 1
    else:
        purge(timestamp+'/pickled_output/','trajectory_data[0-9]*.pkl')
        return 0

def dirs_after_stamp(timestamp):
    dir_list = sorted(os.listdir(os.getcwd()))
    dirs = dir_list[dir_list.index(timestamp):]
    return filter(is_timestamped, dirs)

def bunch_all_output(timestamp):
    for durr in dirs_after_stamp(timestamp):
        err = bunch_output(durr)

def rm_all_pickles(timestamp):
    for durr in dirs_after_stamp(timestamp):
        err = rm_pickles(durr)

def transfer_data_arrays(tmstmp):
    """
    Saves the max value of epsilon and tmeas_off to a dict with the
    output.
    """

    inpt   = '/trajectory_input_file'    
    otpt   = '/output_arrays/output.pkl'
    sv_pr  = 'perm_data/output'
    
    for durr in dirs_after_stamp(tmstmp):
        with open(durr + inpt,'r') as phil:
            inpt_dct = pickle.load(phil)
        with open(durr + otpt,'r') as phil:
            otpt_dct = pickle.load(phil)
        
        print inpt_dct.keys()        
        
        otpt_dct['max_eps'] = np.amax(inpt_dct['epsilon'])
        otpt_dct['tmeas_off'] = inpt_dct['tmeas_off']
        
        for key in SOMETIMES_KEYS:    
            if key in inpt_dct.keys():
                otpt_dct[key] = inpt_dct[key]
        with open(sv_pr + durr + '.pkl', 'w') as phil:
            pickle.dump(otpt_dct, phil)                

if __name__ == "__main__":
    timestamp = argv[1]
    if is_timestamped(timestamp):
        bunch_all_output(timestamp)
    else:
        print "Not a timestamp"
    
